import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans

spr = 2
spc = 2
spn = 0
plt.figure(figsize=[8, 8])

x = np.loadtxt(r'../data/test.txt')

spn += 1
plt.subplot(spr, spc, spn)
plt_x = x[:, 0]
plt_y = x[:, 1]
plt.scatter(plt_x, plt_y, s=1)

spn += 1
plt.subplot(spr, spc, spn)
km = KMeans(4)
km.fit(x)
labels = np.unique(km.labels_)
n_labels = len(labels)
cmap = plt.cm.get_cmap(name='rainbow', lut=n_labels)
plt.scatter(plt_x, plt_y, s=1, c=km.labels_, cmap=cmap, zorder=100)
centers = km.cluster_centers_
for i in labels:
    # color = cmap(n_labels - 1 - i)
    color = cmap(i)
    plt.scatter(centers[i, 0], centers[i, 1], color=color, zorder=0, marker='x')
    plt.annotate('Cluster ' + str(i + 1), xy=[centers[i, 0], centers[i, 1]], color=color, zorder=0)

spn += 1
plt.subplot(spr, spc, spn)
n_k = 10
ks = range(1, n_k+1)
inertias = []
for k in ks:
    km = KMeans(k)
    km.fit(x)
    inertias.append(km.inertia_)
plt.plot(ks, inertias, label='inertia')
for i, k in enumerate(ks):
    plt.annotate(k, xy=[k, inertias[i]])
plt.grid()
plt.legend()

plt.show()
